import matplotlib.pyplot as plt
import numpy as np


def convert_to_grayscale_1(image, channel_weights="rgb" or "equal"):
    "Convert any image into a grayscale image of shape `(h, w)`."

    if isinstance(channel_weights, str):
        if channel_weights == "equal":
            channel_weights = [0.333, 0.334, 0.333]
        if channel_weights == "rgb":
            channel_weights = [0.299, 0.587, 0.114]
    assert isinstance(channel_weights, (list, tuple)) and len(channel_weights) == 3

    source_dtype = image.dtype

    if isinstance(source_dtype, np.floating):
        image = np.asarray(image * 255.0, dtype=np.uint8)
        source_dtype = image.dtype

    if image.ndim == 3:
        image = np.dot(image[..., :3], channel_weights)
        image = np.asarray(image, dtype=source_dtype)

    return image


def convert_to_grayscale_3(image, channel_weights="rgb" or "equal"):
    """Convert any image into a grayscale image of shape `(h, w, 3)`."""

    image = convert_to_grayscale_1(image, channel_weights)
    image = np.repeat(image[..., None], repeats=3, axis=2)

    return image


def load_image_as_grayscale(image_path, channel_weights):
    """Load an image with automatic gray scale conversion.
    Returns a grayscale image as an integer array.
    """

    # The returned array of `plt.imread` has shape:
    #   (M, N) for grayscale images.
    #   (M, N, 3) for RGB images.
    #   (M, N, 4) for RGBA images.
    # PNG images are returned as float arrays (0-1).
    # All other formats are returned as int arrays.
    image = plt.imread(image_path)
    image = convert_to_grayscale_1(image, channel_weights)

    return image


def save_grayscale_image(image_path, image, one_channel=False):
    """Save a grayscale image, expand number of channels to 3 if specified."""

    assert (
        image.ndim == 2
    ), "Designed for grayscale images with 2 channels: (height, width)."

    if not one_channel:
        image = convert_to_grayscale_3(image, "rgb")

    plt.imsave(image_path, image)


__all__ = [
    "convert_to_grayscale_1",
    "convert_to_grayscale_3",
    "load_image_as_grayscale",
    "save_grayscale_image",
]
